/*
 * 环境: ROS Noetic
 * pcd_dynamic_loader.cpp
 * 实现点云的动态读取与发布
 */

#include <ros/ros.h>
#include <pcl_ros/point_cloud.h>
#include <pcl/io/pcd_io.h>
#include <sensor_msgs/PointCloud2.h>
#include <tf/transform_listener.h>
#include <yaml-cpp/yaml.h>
#include <boost/filesystem.hpp>
#include <cmath>

struct Pcd_block_info
{
    std::string path;
    float tf[6]; // x, y, z, roll, pitch, yaw
};

class PCDDynamicLoader
{
public:
    PCDDynamicLoader(ros::NodeHandle &nh, ros::NodeHandle &pnh)
        : nh_(nh), pnh_(pnh), tf_listener_(ros::Duration(10.0))
    {
        // 加载参数
        pnh_.param("tf_check_interval", timer_interval_, 1.0f);
        pnh_.param("tf_from_frame", tf_from_frame_, std::string("map"));
        pnh_.param("tf_to_frame", tf_to_frame_, std::string("body"));
        pnh_.param("load_distance", load_distance_, 2.0f);
        pnh_.param("remove_distance", remove_distance_, 3.0f);
        pnh_.param("pc_publish_topic", pc_publish_topic_, std::string("map3d"));
        pnh_.param("pcd_root_dir", pcd_root_dir_, std::string(""));

        //
        ROS_INFO("<pcd_loader> Reading config from: %s", (pcd_root_dir_ + "pcds_config.yaml").c_str());
        readConfig();

        // 发布器
        pcd_pub_ = nh_.advertise<sensor_msgs::PointCloud2>(pc_publish_topic_, 1, true);

        // 定时器
        timer_ = nh_.createTimer(ros::Duration(timer_interval_), &PCDDynamicLoader::timerCallback, this);

        // 初始化加载原点附近
        float init_tf[6] = {0, 0, 0, 0, 0, 0};
        updateMap(init_tf);
    }

    ~PCDDynamicLoader() {}

private:
    void timerCallback(const ros::TimerEvent &)
    {
        tf::StampedTransform transform;
        try
        {
            tf_listener_.lookupTransform(tf_from_frame_, tf_to_frame_, ros::Time(0), transform);
        }
        catch (tf::TransformException &ex)
        {
            ROS_WARN_THROTTLE(5.0, "TF lookup failed: %s", ex.what());
            return;
        }

        float current_tf[6];
        current_tf[0] = transform.getOrigin().x();
        current_tf[1] = transform.getOrigin().y();
        current_tf[2] = transform.getOrigin().z();
        tf::Matrix3x3 mat(transform.getRotation());
        double roll, pitch, yaw;
        mat.getRPY(roll, pitch, yaw);
        current_tf[3] = roll;
        current_tf[4] = pitch;
        current_tf[5] = yaw;

        // ROS_INFO("Current tf: (%.2f, %.2f, %.2f, %.2f, %.2f, %.2f)", current_tf[0], current_tf[1], current_tf[2],
        //          current_tf[3], current_tf[4], current_tf[5]);
        updateMap(current_tf);
    }

    void updateMap(const float tf_curr[6])
    {
        bool changed = false;
        std::vector<int> to_load;
        // 判断需要加载的块
        for (size_t i = 0; i < pcd_block_info_.size(); ++i)
        {
            if (distance(tf_curr, pcd_block_info_[i].tf) <= load_distance_ &&
                std::find(pcd_block_index_.begin(), pcd_block_index_.end(), i) == pcd_block_index_.end())
            {
                to_load.push_back(i);
            }
        }
        // 加载
        for (int idx : to_load)
        {
            pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
            if (pcl::io::loadPCDFile(pcd_block_info_[idx].path, *cloud) == 0)
            {
                pcd_blocks_.push_back(cloud);
                pcd_block_index_.push_back(idx);
                changed = true;
                ROS_INFO("Loaded block: %s", pcd_block_info_[idx].path.c_str());
            }
            else
            {
                ROS_WARN("Failed to load %s", pcd_block_info_[idx].path.c_str());
            }
        }
        // 删除超过范围的块
        for (size_t i = 0; i < pcd_block_index_.size();)
        {
            int idx = pcd_block_index_[i];
            if (distance(tf_curr, pcd_block_info_[idx].tf) > remove_distance_)
            {
                ROS_INFO("Removing block: %s", pcd_block_info_[idx].path.c_str());
                pcd_blocks_.erase(pcd_blocks_.begin() + i);
                pcd_block_index_.erase(pcd_block_index_.begin() + i);
                changed = true;
            }
            else
            {
                ++i;
            }
        }
        // 重新拼接并发布
        if (changed)
        {
            pcl::PointCloud<pcl::PointXYZ>::Ptr merged(new pcl::PointCloud<pcl::PointXYZ>);
            for (auto &cloud : pcd_blocks_)
            {
                *merged += *cloud;
            }
            sensor_msgs::PointCloud2 output;
            pcl::toROSMsg(*merged, output);
            output.header.frame_id = tf_from_frame_;
            output.header.stamp = ros::Time::now();
            pcd_pub_.publish(output);
            ROS_INFO("Published merged cloud with %zu points", merged->size());
        }
    }

    void readConfig()
    {
        std::string cfg = pcd_root_dir_;
        if (!cfg.empty() && cfg.back() != '/')
            cfg += '/';
        std::string config_file = cfg + "pcds_config.yaml";
        if (!boost::filesystem::exists(config_file))
        {
            ROS_ERROR("Config file not found: %s", config_file.c_str());
            return;
        }
        YAML::Node root = YAML::LoadFile(config_file);
        block_size_ = root["basic"]["block_size"].as<float>();
        // 遍历子点云配置
        for (const auto &node : root["pcds"])
        {
            Pcd_block_info info;
            std::string fname = node["file_name"].as<std::string>();
            info.path = cfg + fname;
            auto tf_arr = node["tf"].as<std::vector<float>>();
            for (int i = 0; i < 6; ++i)
                info.tf[i] = tf_arr[i];
            pcd_block_info_.push_back(info);
        }

        ROS_INFO("<pcd_loader> Block size: %.2f ;Total number of blocks: %zu ", block_size_, pcd_block_info_.size());
    }

    float distance(const float a[6], const float b[6])
    {
        return std::sqrt(std::pow(a[0] - b[0], 2) + std::pow(a[1] - b[1], 2) + std::pow(a[2] - b[2], 2));
    }

    // 成员变量
    ros::NodeHandle nh_, pnh_;
    ros::Timer timer_;
    ros::Publisher pcd_pub_;
    tf::TransformListener tf_listener_;

    // 参数
    float timer_interval_;
    std::string tf_from_frame_, tf_to_frame_;
    float load_distance_, remove_distance_;
    std::string pc_publish_topic_;
    std::string pcd_root_dir_;
    float block_size_;

    std::vector<Pcd_block_info> pcd_block_info_;
    std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> pcd_blocks_;
    std::vector<int> pcd_block_index_;
};

int main(int argc, char **argv)
{
    ros::init(argc, argv, "pcd_dynamic_loader");
    ros::NodeHandle nh;
    ros::NodeHandle pnh("~");
    PCDDynamicLoader loader(nh, pnh);
    ros::spin();
    return 0;
}
